import pymysql
import time
from dbutils.pooled_db import PooledDB


class MySQLConnPool:
    def __init__(self,
                 host='127.0.0.1',
                 port=3306,
                 user='root',
                 password='111111',
                 database='shejiyuan',
                 charset='utf8mb4',
                 maxconnections=10,
                 retry_count=3):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.charset = charset
        self.maxconnections = maxconnections
        self.retry_count = retry_count  # 自动重连的最大次数
        self.pool = None
        self.conn = None
        self.connect_retry_count = 0

    def connect(self):
        """连接MySQL数据库"""
        if not self.pool:
            self.pool = PooledDB(
                creator=pymysql,  # 使用PyMySQL作为连接客户端
                maxconnections=self.maxconnections,
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.password,
                database=self.database,
                charset=self.charset,
                autocommit=False  # 关闭自动提交，手动控制事务
            )
        # 获取数据库连接
        try:
            self.conn = self.pool.connection()
        except Exception as e:
            print(f"Error: {e}")
            if self.connect_retry_count < self.retry_count:  # 自动重连
                self.connect_retry_count += 1
                print(f"Retry connecting to database ({self.connect_retry_count}/{self.retry_count})...")
                time.sleep(1)
                self.connect()
            else:
                print(f"Failed to connect to database ({self.retry_count} attempts have been made).")
                raise e
        else:
            self.connect_retry_count = 0

    def __del__(self):
        """
        析构函数，关闭数据库连接
        """
        if self.conn:
            self.conn.close()

    def close(self):
        """关闭数据库连接"""
        if self.conn:
            self.conn.close()
        self.conn = None

    def commit(self):
        """提交事务"""
        if self.conn:
            self.conn.commit()

    def rollback(self):
        """回滚事务"""
        if self.conn:
            self.conn.rollback()

    def execute_sql(self, sql, args=None):
        """执行SQL语句"""
        cursor = self.conn.cursor()
        try:
            cursor.execute(sql, args)
        except Exception as e:
            print(f"Error: {e}")
            self.rollback()
        else:
            self.commit()
        finally:
            cursor.close()

    def fetchone(self, sql, args=None):
        """查询单条记录"""
        cursor = self.conn.cursor()
        cursor.execute(sql, args)
        result = cursor.fetchone()
        cursor.close()
        return result

    def fetchall(self, sql, args=None):
        """查询多条记录"""
        cursor = self.conn.cursor()
        cursor.execute(sql, args)
        result = cursor.fetchall()
        cursor.close()
        return result

    def dictfetchall(self, sql, args=None):
        """查询多条记录，并返回字典类型的结果"""
        cursor = self.conn.cursor(pymysql.cursors.DictCursor)
        cursor.execute(sql, args)
        result = cursor.fetchall()
        cursor.close()
        return result

    def insert(self, table, data, replace=False, is_bulk=False, batch_size=1000):
        """
        往指定表中插入新记录
        :param table: str 表名
        :param data: dict|list 需要插入的字段及值，键为字段名，值为对应的值。如果 is_bulk 为 True，则 data 是一个包含多个 dict 的列表，每个 dict 代表一条记录
        :param replace: 是否使用replace插入语句，默认为False使用insert语句
        :param is_bulk: bool 是否批量插入，可选参数，默认为 False
        :param batch_size: int 批量插入时的一次性数量
        :return: bool 插入操作是否成功
        """
        sql = "REPLACE" if replace else "INSERT IGNORE"
        if not is_bulk:
            fields = ",".join(data.keys())
            values = tuple(data.values())
            placeholders = ",".join(["%s"] * len(data))
            sql += f" INTO {table} ({fields}) VALUES ({placeholders})"
            try:
                self.execute_sql(sql, values)
                return True
            except Exception as e:
                raise e
        else:
            if len(data) > 0:
                if not data or not isinstance(data, list) or not isinstance(data[0], dict):
                    raise ValueError("The data for bulk insert should be a list of dicts.")
                fields = ",".join(data[0].keys())
                placeholders = ",".join(["%s"] * len(data[0]))
                sql += f" INTO {table} ({fields}) VALUES ({placeholders})"
                try:
                    for i in range(0, len(data), batch_size):
                        data_batch = data[i:i + batch_size]
                        value_list = [tuple(d.values()) for d in data_batch]
                        with self.conn.cursor() as cursor:
                            cursor.executemany(sql, value_list)
                        self.conn.commit()
                    return True
                except Exception as e:
                    self.conn.rollback()
                    raise e

    def select(self, table, fields="*", where=None, order_by=None, desc=False, limit=None, offset=None):
        """
        查询指定表中符合条件的记录
        :param table: 数据表名
        :param fields: 查询字段，默认为查询所有字段（'*'）
        :param where: 查询条件，可以为字符串或者元组类型，默认为None
        :param order_by: 排序方式，可以为字符串或者元组类型，默认为None
        :param desc: 是否降序排列，默认为False（升序排列）
        :param limit: 返回结果的数量上限，可以为整数或者元组类型，默认为10
        :param offset: 偏移量，用于分页，可以为整数，默认为None
        :return: 查询结果，返回list类型的数据
        """
        if isinstance(fields, list):
            fields = ",".join(fields)
        sql = f"SELECT {fields} FROM {table}"
        if where:
            sql += f" WHERE {where}"
        if order_by:
            sql += f" ORDER BY {order_by}"
            sql += " DESC" if desc else " ASC"
        if limit:
            sql += f" LIMIT {limit}"
        if offset:
            sql += f" OFFSET {offset}"
        try:
            result = self.dictfetchall(sql)
            return result
        except Exception as e:
            raise e

